

#include "bamMultiSortCore.h"
#include "utils.h"


bam_hdr_t *h_ ; 


uint64_t num_bps;

int numOutput ;

samFile **fpn  = NULL ;

long *targetOffset_ = NULL ; 

int  numTarget_ =  0 ;


typedef kvec_t(bam1_t) bamVect ;


bamVect  bamVect_ ;


long buckSize =  0 ;

#define BUFFERSIZE 2


void  coreInit(char *fn , char *temp , int  n )
{

    samFile *fp = sam_open(fn,"r");

    fpn =  malloc(sizeof(samFile *)*n);

    kv_resize(bam1_t,bamVect_,BUFFERSIZE);

    kv_setALL(bam1_t,bamVect_);

    h_ = sam_hdr_read(fp);

    int  i = 0 ; 

    numOutput = n ; 

    num_bps = 0 ;

    char *outputName =   NULL;

    outputName = malloc(strlen(temp)+32);

    for( i =  0 ; i < numOutput ; i++){
        sprintf(outputName ,"%s.%d.bam",temp,i);

        fprintf(stderr,"outputName=%s\n",outputName);
        fpn[i] = sam_open(outputName,"wb");
        hts_set_threads(fpn[i],24);
        if(sam_hdr_write(fpn[i],h_) < 0){
            break ;
        }
        memset(outputName , 0 , sizeof(char)*(strlen(temp)+32));
    }

    for ( i = 0 ; i  < h_->n_targets ; i++)
        num_bps += h_->target_len[i];

    buckSize = (int) (((double) (num_bps + 2 ) / (double)n) + 0.5);

    targetOffset_ =  calloc(h_->n_targets,sizeof(long));
    targetOffset_[0] = 0 ; 

    for( i = 1 ;  i  <  h_->n_targets ; i++){
        targetOffset_[i] +=  (h_->target_len[i-1] +  targetOffset_[i-1]);
    } 
    numTarget_ = h_->n_targets ;

    free(outputName);

    sam_close(fp);

}

void core_destory()
{
    bam_hdr_destroy(h_);
}

inline long  offset2coordinate( int tid , int pos)
{
    return  tid >=0 ? targetOffset_[tid] + pos : num_bps ;
}




void  bucketSort()
{
    int i = 0 ;
    for( i = 0 ; i  <  kv_size(bamVect_) ; i++){
        bam1_t  *item =   & bamVect_.a[i] ;
        long offset_ = offset2coordinate(item->core.tid,item->core.pos);
        int  id = offset_ / buckSize ;

//        fprintf(stderr,"id:%d,%ld,%ld\n",id,offset_,buckSize);

        if(sam_write1(fpn[id],h_,item) < 0){
            fprintf(stderr,"Error:Write Bam\n");
            break ;
        }
    } 
    kv_clear(bamVect_);
}


void readBam2Vector( char *fn , int n_threads )
{

    samFile *fp =  sam_open(fn,"r");

    hts_set_threads(fp,n_threads);

    bam_hdr_t *h =  sam_hdr_read(fp);

    long  n = 0 ;

    do{

        bam1_t  *bam =  &bamVect_.a[bamVect_.n] ;
        if( sam_read1(fp,h,bam) < 0 ) {
            break ; 
        }
        
        bamVect_.n++ ;

        n++ ;

        if( bamVect_.m == bamVect_.n ){
            bamVect_.m = bamVect_.m << 1 ;
            bamVect_.a =  realloc(bamVect_.a ,  sizeof(bam1_t)*bamVect_.m);
            memset(bamVect_.a+bamVect_.n,0,sizeof(bam1_t)*(bamVect_.m-bamVect_.n));
        }

        if( kv_size(bamVect_) ==  BUFFERSIZE ){
            bucketSort();
        }

    }while(1);

    if( kv_size(bamVect_) != 0 ){
        bucketSort();
    }

    fprintf(stderr,"Read Bams Numbers %ld\n",n);
    //remove(fn);
    bam_hdr_destroy(h);
    sam_close(fp);
}

static void  bucketReadSort( char *fn , int n_threads )
{
    samFile *fp =  sam_open(fn,"r");

    hts_set_threads(fp,n_threads);

    bam_hdr_t *h =  sam_hdr_read(fp);


    double time1 =  get_run_time();


    kv_clear(bamVect_);

    do{

        bam1_t  *bam =  &bamVect_.a[bamVect_.n] ;
        if( sam_read1(fp,h,bam) < 0 ) {
            break ; 
        }
        bamVect_.n++ ;
        if( bamVect_.m == bamVect_.n ){
            bamVect_.m = bamVect_.m << 1 ;
            bamVect_.a =  realloc(bamVect_.a ,  sizeof(bam1_t)*bamVect_.m);
            memset(bamVect_.a+bamVect_.n,0,sizeof(bam1_t)*(bamVect_.m-bamVect_.n));
        }
    }while(1);

    fprintf(stderr,"Read File %s  = %zu :cost Time =  %2.f", fn,bamVect_.n,get_run_time() - time1);

    bam_hdr_destroy(h);
    sam_close(fp);
    remove(fn);

}

int  compared(const void *a_ ,const void *b_ )
{

    const bam1_t *a = a_ ;
    const bam1_t *b = b_ ;
    return  (((uint64_t)a->core.tid<<32|(a->core.pos+1)<<1|bam_is_rev(a)) > ((uint64_t)b->core.tid<<32|(b->core.pos+1)<<1|bam_is_rev(b))); 
}

void  sortVect()
{
    qsort(bamVect_.a,kv_size(bamVect_),sizeof(bam1_t),compared);
}


void ouputBam(samFile *fp)
{

    int  i = 0 ;

    for( i = 0 ; i <  kv_size(bamVect_) ; i++ ){
        bam1_t *item =  &kv_A(bamVect_,i);
        if(sam_write1(fp,h_,item) < 0){
            fprintf(stderr,"[Error] Write Bams\n");
            break ;
        }
    }
}


void sortCore_( char *inputName ,samFile *fp,  int n_threads )
{

    bucketReadSort( inputName  , n_threads );
    // sort //
    double time1 =  get_run_time();
    sortVect();
    fprintf(stderr," Sort fun Time = %2.f\n", get_run_time() - time1);
    time1 = get_run_time();
    // outout 
    ouputBam(fp);
    fprintf(stderr," Output fun Time = %2.f , size = %zu\n", get_run_time() - time1 , kv_size(bamVect_));
    
}

void  Runsort( char  *outputBam ,  char * temp ,int n_threads  )
{

    char *inputName =   NULL;

    samFile *fp =  sam_open(outputBam ,"wb");

    hts_set_threads(fp,n_threads);

    // First : close  sam handler
    int  i = 0 ;

    for( i = 0 ; i <  numOutput ; i++ )
        sam_close(fpn[i]);

    free(fpn);


    if( sam_hdr_write(fp,h_) < 0 ){
        fprintf(stderr,"[Error] write head\n");
    }

    inputName = malloc(strlen(temp)+32);

    for( i =  0 ; i < numOutput ; i++){
        sprintf(inputName,"%s.%d.bam",temp,i);
        fprintf(stderr,"inputName =%s\n",inputName );
        sortCore_(inputName , fp , n_threads );
        memset(inputName , 0 , sizeof(char)*(strlen(temp)+32));
    }
    //
    sam_close(fp);
}








